import pandas as pd
from faker import Faker
import random
import os

# 初始化Faker
fake = Faker('zh_CN')

# --- 创建警员信息表 ---
police_data = []
police_ids = [f'P{202500 + i:04d}' for i in range(20)]
for pid in police_ids:
    police_data.append({
        'police_id': pid,
        'name': fake.name(),
        'gender': random.choice(['男', '女']),
        'rank': random.choice(['一级警司', '二级警司', '三级警司', '一级警员', '二级警员'])
    })
police_df = pd.DataFrame(police_data)

# --- 创建案件记录表 ---
case_data = []
alarm_types = ['盗窃', '诈骗', '抢劫', '故意伤害', '网络犯罪', '其他']
locations = ['中山路', '人民路', '解放路', '北京西路', '广州路', '珠江路']
case_ids = [3025000 + i for i in range(100)]

for cid in case_ids:
    case_data.append({
        'case_id': cid,
        'report_time': fake.date_time_between(start_date='-2y', end_date='now'),
        'case_type': random.choice(alarm_types),
        'location': random.choice(locations) + str(random.randint(1, 200)) + '号',
        'handler_id': random.choice(police_ids) # 确保办案警员ID存在
    })
case_df = pd.DataFrame(case_data)

# --- 创建嫌疑人信息表 ---
suspect_data = []
for _ in range(50):
    suspect_data.append({
        'suspect_id': fake.ssn(),
        'name': fake.name(),
        'age': random.randint(18, 60),
        'involved_case_id': random.choice(case_ids) # 确保涉案ID存在
    })
suspect_df = pd.DataFrame(suspect_data)

# --- 保存为CSV文件 ---
output_dir = 'final_product'
os.makedirs(output_dir, exist_ok=True)

police_df.to_csv(f'{output_dir}/mock_police.csv', index=False, encoding='utf-8-sig')
case_df.to_csv(f'{output_dir}/mock_cases.csv', index=False, encoding='utf-8-sig')
suspect_df.to_csv(f'{output_dir}/mock_suspects.csv', index=False, encoding='utf-8-sig')

print(f"数据库模拟数据已成功生成到 {output_dir} 目录下：mock_police.csv, mock_cases.csv, mock_suspects.csv")
